Training Communication-Efficient Mixture-of-Experts Language Models with Layer Re-Configuration
训练通信高效 Mixture-of-Experts 语言模型的层重配置方法
🧭 通信高效 MoE 架构训练优化 ★ 深度解读
NVIDIA
减少 MoE 层数并用 Mamba-2 和密集 FFN 补深度，在匹配参数下降低专家并行通信和训练 GPU 小时 30-35%
📌 工作参考价值（4 点）
- 如果业务里用 MoE LLM 做召回/粗排/reranker 或生成式推荐，可尝试把 MoE 层从每个 token-mixing 后都放一个改为稀疏放置，例如 52 层 hybrid 模型只保留 8 个 MoE 层，其余用 Mamba-2 和 dense FFN 补深度，理论上 all-to-all 调用从 23q 降到 8q，实际训练 GPU 小时省 30% 以上。
- 层排序不要随机，采用 phased-greedy：先把 attention 和 channel-mixing 块按虚拟块分配到 Mamba spine，避免连续 channel-mixing，attention 放在 token-mixing 尾部；实验中比 21 个受约束随机排序的验证损失低约 1.8σ。
- 连续堆叠 softmax attention 容易 rank collapse 导致训练发散；如果用 token-mixing-heavy 架构，优先用 Mamba-2/线性注意力作为主干，比 softmax attention 稳定得多。
- 评估架构改动要直接看 GPU-hours 而非理论 FLOPs/参数量；本文的模型形状搜索用 L_E*K 作为通信代价代理，在参数匹配候选中选该值最小者，可复用到自己的搜索流程。